package mr

import (
	"fmt"
	"io/ioutil"
	"os"
	"sort"
	"strings"
	"time"
)
import "log"
import "net/rpc"
import "hash/fnv"

//
// Map functions return a slice of KeyValue.
//
type KeyValue struct {
	Key   string
	Value string
}

//
// use ihash(key) % NReduce to choose the reduce
// task number for each KeyValue emitted by Map.
//
func ihash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

func contains(a []string, x string) bool {
	for _, n := range a {
		if x == n {
			return true
		}
	}
	return false
}

func mapKeys(m map[string][]string) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	return keys
}

func groupByKey(kva []KeyValue) map[string][]KeyValue {
	dict := make(map[string][]KeyValue)
	for _, kv := range kva {
		if dict[kv.Key] != nil {
			dict[kv.Key] = append(dict[kv.Key], kv)
		} else {
			dict[kv.Key] = []KeyValue{kv}
		}
	}
	return dict
}

func partitionKva(kva []KeyValue, taskId int) []string {
	filenames := []string{}
	groups := groupByKey(kva)
	for key, values := range groups {
		// hard code for nReduce
		n := ihash(key) % 10
		_, err2 := os.Stat("out")
		if os.IsNotExist(err2) {
			os.Mkdir("out", os.ModePerm)
		}
		outFileName := fmt.Sprintf("out/m-out-%d-t%d", n, taskId)
		_, err := os.Stat(outFileName)
		var file *os.File
		if os.IsNotExist(err) {
			outFile, err := os.Create(outFileName)
			if err != nil {
				log.Printf("file %s can not create", outFileName)
				continue
			}
			file = outFile
		} else {
			file, err = os.OpenFile(outFileName, os.O_APPEND|os.O_WRONLY, 0666)
			if err != nil {
				log.Printf("file %s can not open", outFileName)
				continue
			}
		}
		for _, value := range values {
			file.WriteString(fmt.Sprintf("%v %v\n", value.Key, value.Value))
		}
		file.Close()
		if !contains(filenames, outFileName) {
			filenames = append(filenames, outFileName)
		}
	}
	return filenames
}

//
// main/mrworker.go calls this function.
//
func Worker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {

	// Your worker implementation here.

	// uncomment to send the Example RPC to the master.
	//CallExample()

	for {
		reply := CallGetTask()
		if reply != nil && reply.DidGetTask {
			if reply.TaskType == "map" {
				MapWorker(reply, mapf)
			} else if reply.TaskType == "reduce" {
				ReduceWorker(reply, reducef)
			}
		}
		time.Sleep(time.Millisecond * 500)
	}

}

// Worker for map
func MapWorker(reply *GetTaskReply, mapf func(string, string) []KeyValue) {
	if reply != nil && reply.DidGetTask {
		file, err := os.Open(reply.FilePath)
		if err != nil {
			log.Fatalf("cannot open %v", reply.FilePath)
		}
		content, err := ioutil.ReadAll(file)
		if err != nil {
			log.Fatalf("cannot read %v", reply.FilePath)
		}
		file.Close()
		// call map function
		kva := mapf(reply.FilePath, string(content))
		// partition
		filenames := partitionKva(kva, reply.TaskId)
		// send reply
		ret := FinishTaskArgs{
			TaskId:    reply.TaskId,
			Filenames: filenames,
		}
		CallFinishTask(&ret)
	}
}

// Worker for reduce
func ReduceWorker(reply *GetTaskReply, reducef func(string, []string) string) {
	if reply != nil && reply.DidGetTask {
		contentMap := make(map[string][]string)
		// read split files
		for _, filePath := range reply.FilePaths {
			file, err := os.Open(filePath)
			if err != nil {
				log.Fatalf("cannot open %v", filePath)
			}
			content, err := ioutil.ReadAll(file)
			if err != nil {
				log.Fatalf("cannot read %v", filePath)
			}
			file.Close()
			lines := strings.Split(string(content), "\n")
			for _, line := range lines {
				if len(line) == 0 || line == " " {
					continue
				}
				row := strings.Split(line, " ")
				key := row[0]
				value := row[1]
				if contentMap[key] != nil {
					contentMap[key] = append(contentMap[key], value)
				} else {
					contentMap[key] = []string{value}
				}
			}
		}
		// sort map keys
		keys := mapKeys(contentMap)
		sort.Strings(keys)
		// result file for reduce task
		outFileName := fmt.Sprintf("mr-out-%d", reply.TaskId)
		outFile, err := os.Create(outFileName)
		if err != nil {
			log.Fatalf("cannot create %v", outFileName)
		}
		for _, key := range keys {
			// call reduce function for each key
			result := reducef(key, contentMap[key])
			outFile.WriteString(fmt.Sprintf("%v %v\n", key, result))
		}
		outFile.Close()

		// send reply
		ret := FinishTaskArgs{
			TaskId:   reply.TaskId,
			Filename: outFileName,
		}
		CallFinishTask(&ret)
	}
}

//
// example function to show how to make an RPC call to the master.
//
// the RPC argument and reply types are defined in rpc.go.
//
func CallExample() {

	// declare an argument structure.
	args := ExampleArgs{}

	// fill in the argument(s).
	args.X = 99

	// declare a reply structure.
	reply := ExampleReply{}

	// send the RPC request, wait for the reply.
	call("Master.Example", &args, &reply)

	// reply.Y should be 100.
	fmt.Printf("reply.Y %v\n", reply.Y)
}

func CallGetTask() *GetTaskReply {

	// declare an argument structure.
	args := GetTaskArgs{}

	// declare a reply structure.
	reply := GetTaskReply{
		DidGetTask: false,
	}

	// send the RPC request, wait for the reply.
	if call("Master.GetTask", &args, &reply) {
		if reply.DidGetTask {
			filenames := ""
			if reply.TaskType == "map" {
				filenames = reply.FilePath
			} else {
				filenames = strings.Join(reply.FilePaths, ",")
			}
			fmt.Printf("get a %s task: %d, path: %s\n", reply.TaskType, reply.TaskId, filenames)
			// reply
			return &reply
		}
	}

	return nil
}

func CallFinishTask(args *FinishTaskArgs) *FinishTaskReply {

	// declare a reply structure.
	reply := FinishTaskReply{}

	// send the RPC request, wait for the reply.
	if call("Master.FinishTask", &args, &reply) {
		return &reply
	}
	return nil
}

//
// send an RPC request to the master, wait for the response.
// usually returns true.
// returns false if something goes wrong.
//
func call(rpcname string, args interface{}, reply interface{}) bool {
	// c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
	sockname := masterSock()
	c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	err = c.Call(rpcname, args, reply)
	if err == nil {
		return true
	}

	fmt.Println(err)
	return false
}
